{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3828b987",
   "metadata": {},
   "source": [
    "# 训练流程\n",
    "## 前言\n",
    "本节是很重要的一节，我们将学习Model接口这个高阶API的基本用法，以及Model接口内部的中低阶API的使用，也就是解剖Model接口。在这些基础之上，学习如何自定义这些中低阶API来满足多元的训练场景。\n",
    "Model接口部分功能如下:\n",
    "![](./Model.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "467a50ba",
   "metadata": {},
   "source": [
    "Model接口具有连接神经网络、损失函数构成`损失网络或者评估网络`，连接损失网络和优化器构成`训练网络`，以及`执行训练、评估和推理`的功能。我们把Model接口拆开来看，其实是对一些封装好的中低阶API的再次封装。\n",
    "\n",
    "- 连接神经网络和损失函数组成损失网络: `nn.WithLossCell(net, loss_fn)`\n",
    "- 连接神经网络和损失函数组成评估网络: `nn.WithEvalCell(net, loss_fn)`\n",
    "- 连接损失网络和优化器: `nn.TrainOneStepCell(net_with_loss, optim)`\n",
    "- 执行训练: `net.set_train()`\n",
    "- 执行评估: `net.set_train(False)`\n",
    "- 使用`for循环`进行迭代"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d7c928d",
   "metadata": {},
   "source": [
    "我们也可以针对不同的任务场景，自定义上述网络。值得一提的是，上述网络是MindSpore函数式编程下独有的。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23c7a317",
   "metadata": {},
   "source": [
    "## 一、Model接口的使用"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7cc55ef",
   "metadata": {},
   "source": [
    "`Model`是MindSpore提供的高阶API，可以进行模型训练、评估和推理。其接口的常用参数如下：\n",
    "\n",
    "- `network`：用于训练或推理的神经网络。\n",
    "- `loss_fn`：所使用的损失函数。\n",
    "- `optimizer`：所使用的优化器。\n",
    "- `metrics`：用于模型评估的评价函数。\n",
    "- `eval_network`：模型评估所使用的网络，未定义情况下，`Model`会使用`network`和`loss_fn`进行封装。\n",
    "\n",
    "`Model`提供了以下接口用于模型训练、评估和推理：\n",
    "\n",
    "- `train`：用于在训练集上进行模型训练。\n",
    "- `eval`：用于在验证集上进行模型评估。\n",
    "- `predict`：用于对输入的一组数据进行推理，输出预测结果。\n",
    "\n",
    "对于简单场景的神经网络，可以在定义`Model`时指定前向网络`network`、损失函数`loss_fn`、优化器`optimizer`和评价函数`metrics`。\n",
    "\n",
    "此时，`Model`会使用`network`作为前向网络，并使用`nn.WithLossCell`和`nn.TrainOneStepCell`构建训练网络，使用`nn.WithEvalCell`构建评估网络。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13e2500d",
   "metadata": {},
   "source": [
    "我们依旧以简单线性为例，直接薅过来之前的代码："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "c295239f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:03.747181Z",
     "start_time": "2022-11-13T07:20:03.712338Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(34264:34356,MainProcess):2022-11-13-15:20:03.723.980 [mindspore\\dataset\\engine\\datasets_user_defined.py:656] Python multiprocessing is not supported on Windows platform.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import mindspore.dataset as ds\n",
    "import mindspore.nn as nn\n",
    "import mindspore as ms\n",
    "from mindspore.common.initializer import Normal\n",
    "from mindvision.engine.callback import LossMonitor\n",
    "\n",
    "def get_data(num, w=2.0, b=3.0):\n",
    "    \"\"\"生成样本数据及对应的标签\"\"\"\n",
    "    for _ in range(num):\n",
    "        x = np.random.uniform(-10.0, 10.0)\n",
    "        noise = np.random.normal(0, 1)\n",
    "        y = x * w + b + noise\n",
    "        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)\n",
    "\n",
    "def create_dataset(num_data, batch_size=16):\n",
    "    \"\"\"生成数据集\"\"\"\n",
    "    dataset = ds.GeneratorDataset(list(get_data(num_data)), column_names=['data', 'label'])\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    return dataset\n",
    "\n",
    "class LinearNet(nn.Cell):\n",
    "    \"\"\"定义线性回归网络\"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Dense(1, 1, Normal(0.02), Normal(0.02))\n",
    "\n",
    "    def construct(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "train_dataset = create_dataset(num_data=160)\n",
    "net = LinearNet()\n",
    "loss_fn = nn.MSELoss()\n",
    "optim = nn.Momentum(net.trainable_params(), learning_rate=0.005, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b7b1d8",
   "metadata": {},
   "source": [
    "初识化`Model`,构建训练网络："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "0295cf93",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:03.763255Z",
     "start_time": "2022-11-13T07:20:03.749194Z"
    }
   },
   "outputs": [],
   "source": [
    "model = ms.Model(network=net, loss_fn=loss_fn, optimizer=optim, metrics={\"mae\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e25059c",
   "metadata": {},
   "source": [
    "执行训练："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "7b6eee3f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:03.936754Z",
     "start_time": "2022-11-13T07:20:03.766374Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:[  0/  1], step:[    1/   10], loss:[167.153/167.153], time:119.000 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    2/   10], loss:[60.065/113.609], time:1.987 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    3/   10], loss:[10.604/79.274], time:1.006 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    4/   10], loss:[43.743/70.391], time:2.067 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    5/   10], loss:[96.826/75.678], time:0.998 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    6/   10], loss:[81.910/76.717], time:2.927 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    7/   10], loss:[26.842/69.592], time:0.999 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    8/   10], loss:[9.569/62.089], time:0.940 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[    9/   10], loss:[56.498/61.468], time:1.004 ms, lr:0.00500\n",
      "Epoch:[  0/  1], step:[   10/   10], loss:[51.890/60.510], time:1.004 ms, lr:0.00500\n",
      "Epoch time: 150.002 ms, per step time: 15.000 ms, avg loss: 60.510\n"
     ]
    }
   ],
   "source": [
    "model.train(epoch=1, train_dataset=train_dataset, callbacks=[LossMonitor(0.005)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "856e32d3",
   "metadata": {},
   "source": [
    "模型评估："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "5420acad",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:03.998696Z",
     "start_time": "2022-11-13T07:20:03.941692Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(34264:34356,MainProcess):2022-11-13-15:20:03.945.690 [mindspore\\dataset\\engine\\datasets_user_defined.py:656] Python multiprocessing is not supported on Windows platform.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'mae': 5.784346961975098}\n"
     ]
    }
   ],
   "source": [
    "eval_dataset = create_dataset(num_data=80)\n",
    "mae = model.eval(eval_dataset)\n",
    "print(mae)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d4b92a",
   "metadata": {},
   "source": [
    "## 二、解剖Model接口\n",
    "上面已经提到，`Model`接口其实就是封装了一些连接神经网络、损失函数和优化器的API，现在我们来分别看看这些中低阶API："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443e08fd",
   "metadata": {},
   "source": [
    "我们再复制粘贴一下之前的代码："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e55e33",
   "metadata": {},
   "source": [
    "### nn.WithLossCell\n",
    "`WithLossCell`这个接口可以连接前向网络和损失函数，组成损失网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "d0c6e61c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:04.014685Z",
     "start_time": "2022-11-13T07:20:04.000654Z"
    }
   },
   "outputs": [],
   "source": [
    "net = LinearNet()\n",
    "loss = nn.MSELoss()\n",
    "\n",
    "# 创建损失网络\n",
    "net_with_loss = nn.WithLossCell(net, loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a70228c7",
   "metadata": {},
   "source": [
    "### nn.TrainOneStepCell\n",
    "`TrainOneStepCell`这个接口可以连接损失网络和优化器，组成训练网络。\n",
    "参数如下：\n",
    "\n",
    "- `network`：训练网络，只支持单输出网络。\n",
    "- `optimizer`： 用于更新网络参数的优化器。\n",
    "- `sens`：反向传播的输入，缩放系数，默认值为1.0。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "581353b2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:04.030257Z",
     "start_time": "2022-11-13T07:20:04.016659Z"
    }
   },
   "outputs": [],
   "source": [
    "# 创建训练网络\n",
    "optim = nn.Momentum(net.trainable_params(), learning_rate=0.005, momentum=0.9)\n",
    "train_net = nn.TrainOneStepCell(net_with_loss, optim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2fad5ab",
   "metadata": {},
   "source": [
    "使用`set_train`通过`mode`参数指定模型是否为训练模式，其中`mode`参数默认为True，即默认情况下为训练模式，若`mode`为False，则为评估或推理模式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "e6ff18c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:04.205415Z",
     "start_time": "2022-11-13T07:20:04.032267Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(34264:34356,MainProcess):2022-11-13-15:20:04.402.65 [mindspore\\dataset\\engine\\datasets_user_defined.py:656] Python multiprocessing is not supported on Windows platform.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [0 / 2], step: [1 / 10], loss: 180.86636\n",
      "Epoch: [0 / 2], step: [2 / 10], loss: 61.35399\n",
      "Epoch: [0 / 2], step: [3 / 10], loss: 11.327038\n",
      "Epoch: [0 / 2], step: [4 / 10], loss: 30.328594\n",
      "Epoch: [0 / 2], step: [5 / 10], loss: 109.6683\n",
      "Epoch: [0 / 2], step: [6 / 10], loss: 115.98832\n",
      "Epoch: [0 / 2], step: [7 / 10], loss: 25.788277\n",
      "Epoch: [0 / 2], step: [8 / 10], loss: 6.741369\n",
      "Epoch: [0 / 2], step: [9 / 10], loss: 16.145555\n",
      "Epoch: [0 / 2], step: [10 / 10], loss: 46.796722\n",
      "Epoch: [1 / 2], step: [1 / 10], loss: 40.311687\n",
      "Epoch: [1 / 2], step: [2 / 10], loss: 21.276863\n",
      "Epoch: [1 / 2], step: [3 / 10], loss: 4.1802406\n",
      "Epoch: [1 / 2], step: [4 / 10], loss: 3.5435543\n",
      "Epoch: [1 / 2], step: [5 / 10], loss: 10.153761\n",
      "Epoch: [1 / 2], step: [6 / 10], loss: 27.21175\n",
      "Epoch: [1 / 2], step: [7 / 10], loss: 12.518449\n",
      "Epoch: [1 / 2], step: [8 / 10], loss: 20.859024\n",
      "Epoch: [1 / 2], step: [9 / 10], loss: 1.5066166\n",
      "Epoch: [1 / 2], step: [10 / 10], loss: 7.4648147\n"
     ]
    }
   ],
   "source": [
    "# 生成训练数据集\n",
    "train_dataset = create_dataset(num_data=160, batch_size=16)\n",
    "\n",
    "# 执行训练\n",
    "train_net.set_train()\n",
    "\n",
    "step = 1\n",
    "epochs = 2\n",
    "steps = train_dataset.get_dataset_size()\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for d in train_dataset.create_dict_iterator():\n",
    "        result = train_net(d[\"data\"], d[\"label\"])\n",
    "        print(f\"Epoch: [{epoch} / {epochs}], \"\n",
    "              f\"step: [{step} / {steps}], \"\n",
    "              f\"loss: {result}\")\n",
    "        step = step + 1\n",
    "        if step == 11:\n",
    "            step = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ae84ac",
   "metadata": {},
   "source": [
    "### nn.WithEvalCell\n",
    "MindSpore的`nn`模块提供了评估网络封装函数`WithEvalCell`，用来在验证集上评估模型训练的效果。其参数如下：\n",
    "\n",
    "- `network`：前向网络。\n",
    "- `loss_fn`：损失函数。\n",
    "- `add_cast_fp32`：是否将数据类型调整为float32。\n",
    "\n",
    "`nn.WithEvalCell`只接受两个输入，分别为数据`data`及其对应的标签`label`，用前面定义的前向网络和损失函数构建一个评估网络，示例如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "06fe50bf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T07:20:04.284249Z",
     "start_time": "2022-11-13T07:20:04.208414Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(34264:34356,MainProcess):2022-11-13-15:20:04.215.416 [mindspore\\dataset\\engine\\datasets_user_defined.py:656] Python multiprocessing is not supported on Windows platform.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mae:  3.441567134857178\n",
      "loss:  16.285030460357667\n"
     ]
    }
   ],
   "source": [
    "eval_dataset = create_dataset(num_data=160, batch_size=16)\n",
    "\n",
    "# 创建评估网络\n",
    "eval_net = nn.WithEvalCell(net, loss)\n",
    "eval_net.set_train(False)\n",
    "\n",
    "loss = nn.Loss() # 这个方法是计算Loss的平均值，可查看官方文档\n",
    "mae = nn.MAE()\n",
    "\n",
    "# 初始化两个方法\n",
    "mae.clear()\n",
    "loss.clear()\n",
    "\n",
    "# 真正验证迭代过程\n",
    "for data in eval_dataset.create_dict_iterator():\n",
    "    outputs = eval_net(data[\"data\"], data[\"label\"])\n",
    "    mae.update(outputs[1], outputs[2])\n",
    "    loss.update(outputs[0])\n",
    "\n",
    "# 评估结果\n",
    "mae_result = mae.eval()\n",
    "loss_result = loss.eval()\n",
    "\n",
    "print(\"mae: \", mae_result)\n",
    "print(\"loss: \", loss_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "196b2878",
   "metadata": {},
   "source": [
    "## 三、自定义\n",
    "针对不同的任务场景，MindSpore内置的API可能无法满足我们的需求，比如当多标签的情况下，下面我们基于多标签任务，学习自定义损失网络、训练网络和评估网络。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82bf8124",
   "metadata": {},
   "source": [
    "### 自定义损失网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "6e531d65",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T08:14:06.368610Z",
     "start_time": "2022-11-13T08:14:06.343094Z"
    }
   },
   "outputs": [],
   "source": [
    "class CustomWithLossCell(nn.Cell):\n",
    "    \"\"\"自定义WithLossCell接口\"\"\"\n",
    "    def __init__(self, backbone, loss_fn):\n",
    "        # 初始化网络和损失函数\n",
    "        super(CustomWithLossCell, self).__init__(auto_prefix=False)\n",
    "        self.backbone = backbone\n",
    "        self.loss_fn = loss_fn\n",
    "        \n",
    "    def construct(self, data, label1, label2):\n",
    "        # 构建网络\n",
    "        x = self.backbone(data) # 网络预测值\n",
    "        loss = self.loss_fn(x, label1, label2) # 损失值\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "984c4aca",
   "metadata": {},
   "source": [
    "### 自定义训练网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "0a76c691",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T08:11:08.095920Z",
     "start_time": "2022-11-13T08:11:08.079083Z"
    }
   },
   "outputs": [],
   "source": [
    "class CustomTrainOneStepCell(nn.Cell):\n",
    "    \"\"\"自定义TrainOneStepCell接口\"\"\"\n",
    "    def __init__(self, net, optim):\n",
    "        # 初始化损失网络和优化器\n",
    "        super(CustomTrainOneStepCell, self).__init__(auto_prefix=False)\n",
    "        self.net = net                                  # 初始化前向网络\n",
    "        self.net.set_grad()                             # 构建反向网络\n",
    "        self.optim = optim                              # 初始化优化器\n",
    "        self.weights = self.optim.parameters            # 待更新的参数  \n",
    "        self.grad = ops.GradOperation(get_by_list=True) # 反向传播获取梯度\n",
    "    \n",
    "    def construct(*inputs):\n",
    "        loss = self.net(*inputs)                           # 计算当前loss\n",
    "        grads = self.grad(self.net, self.weights)(*inputs) # 反向传播计算梯度\n",
    "        self.optim(grads)                                  # 使用优化器更新权重\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe2b30e8",
   "metadata": {},
   "source": [
    "### 自定义评估网络\n",
    "在自定义时，如不需要损失函数作为评价指标，则无需定义`loss_fn`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "fb813d58",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-13T08:15:45.714482Z",
     "start_time": "2022-11-13T08:15:45.694486Z"
    }
   },
   "outputs": [],
   "source": [
    "class CustomWithEvalCell(nn.Cell):\n",
    "    \"\"\"自定义CustomWithEvalCell接口\"\"\"\n",
    "    def __init__(self, net):\n",
    "        super(CustomWithEvalCell, self).__init__(auto_prefix=False)\n",
    "        self.net = net\n",
    "        \n",
    "    def construct(self, data, label1, label2):\n",
    "        outputs = self.net(data)\n",
    "        return outpus, label1, label2"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore",
   "language": "python",
   "name": "mindvision"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
